#!/usr/bin/env python3
# Author: Armit
# Create Time: 2023/10/09

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor


class ANN(nn.Module):

  def __init__(self, dim):
    super(ANN,self).__init__()

    self.linear1  = nn.Linear(dim, 512) ; self.norm1 = nn.LayerNorm(512)
    self.linear2  = nn.Linear(512, 256) ; self.norm2 = nn.LayerNorm(256)
    self.linear3  = nn.Linear(256, 128) ; self.norm3 = nn.LayerNorm(128)
    self.linear4  = nn.Linear(128, 96)  ; self.norm4 = nn.LayerNorm(96)
    self.linear5  = nn.Linear(96, 64)
    self.linear6  = nn.Linear(64, 32)
    self.linear7  = nn.Linear(32, 16)
    self.linear8  = nn.Linear(16, 8)
    self.linear9  = nn.Linear(8, 4)
    self.linear10 = nn.Linear(4, 2)

  def forward(self, x:Tensor) -> Tensor:
    x = F.relu(self.norm1(self.linear1(x)))
    x = F.relu(self.norm2(self.linear2(x)))
    x = F.relu(self.norm3(self.linear3(x)))
    x = F.relu(self.norm4(self.linear4(x)))
    x = F.relu(self.linear5(x))
    x = F.relu(self.linear6(x))
    x = F.relu(self.linear7(x))
    x = F.relu(self.linear8(x))
    x = F.relu(self.linear9(x))
    x: Tensor = self.linear10(x)
    x = x.squeeze(-1)
    return x


class Myloss(nn.Module):

  def __init__(self, alpha, beta):
    super(Myloss,self).__init__()

    self.alpha = alpha
    self.beta = beta

  def forward(self, g, g_hat, a, a_hat) -> Tensor:
    return torch.mean(self.alpha * torch.pow((g - g_hat), 2) + self.beta * torch.pow((a - a_hat), 2))


def r2(output:Tensor, target:Tensor) -> Tensor:
  target_mean = torch.mean(target)
  ss_tot = torch.sum((target-target_mean)**2)
  ss_res = torch.sum((target-output)**2)
  r2_loss = 1 - ss_res / ss_tot
  return r2_loss
